package cn.edu.sjtu;

import hu.uw.pallergabor.dedexer.Dedexer;
import hu.uw.pallergabor.dedexer.JasminStyleCodeGenerator;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileReader;
import java.io.InputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;

import org.w3c.dom.Document;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import cn.edu.sjtu.intent.IntentChecker;
import cn.edu.sjtu.permission.APIPermissionChecker;
import cn.edu.sjtu.permission.CPPermissionChecker;
import cn.edu.sjtu.permission.IntentPermissionChecker;
import cn.edu.sjtu.permission.PermissionChecker;
import cn.edu.sjtu.permissionleak.InvokeCompleted;
import cn.edu.sjtu.todo.Cfg;
import cn.edu.sjtu.todo.nodes.MethodNode;

public class DedexerPlus {

	public static HashMap<String, HashMap<Long, List<HashMap<String, String>>>> registerMapsForClasses;
    public final static HashMap<String, String> androidMap = new HashMap<String, String>();	// for android class's fields
	public static HashMap<Long, Long> filePos2Line;
	public static HashMap<Long, List<Long>> jumpMap;
	public static HashMap<String, String> extendMap;
	public static HashMap<String, List<String>> reExtendMap;
	
	private HashSet<String> permission_app_required = null;
	private HashSet<String> permission_app_acquired = null;
	
	private HashSet<String> permission_func = null;
	private static HashMap<String, HashSet<String>> permission_func_required = null;
	
	public static HashMap<String, List<String>> invokeMap;
	public static HashMap<String, List<String>> reInvokeMap;
	
	private String rootPath = null;
	
	private PermissionChecker apipc = null;
	private PermissionChecker intpc = null;
	private PermissionChecker cppc = null;
	private IntentChecker intc = null;
	private PermissionChecker ic = null;
	
	private Cfg cb = null;
	
	public static boolean isWriteDDX;
	public static boolean isPermissionCheck = true;
	public static boolean isIntentCheck = false;
	public static boolean isInvokeCompleted = false;
	public static boolean isCfgBuild = false;

	public static long totalMethod = 0;
	public static long totalMethod_ = 0;
	
	public DedexerPlus(String cpXlsFilePath, String apiXlsFilePath, String intXlsFilePath, String am) {	
		apipc = new APIPermissionChecker(apiXlsFilePath);
		intpc = new IntentPermissionChecker(intXlsFilePath);
		cppc = new CPPermissionChecker(cpXlsFilePath);
		if (isIntentCheck) {
			intc = new IntentChecker();
		}
		if (isInvokeCompleted) {
			ic = new InvokeCompleted(null);
			permission_func_required = new HashMap<String, HashSet<String>>();
			permission_func = new HashSet<String>();
		}
		
		if (isInvokeCompleted) {
			invokeMap = new HashMap<String, List<String>>();
			reInvokeMap = new HashMap<String, List<String>>();
		}
		
		if (isCfgBuild) {
			cb = new Cfg();
		}
		
		initAndroidMap(am);
	}
	
	public static String getFatherClass(String c) {
		if (extendMap != null) {
			return extendMap.get(c);
		}
		return null;
	}
	
	private void initAndroidMap(String am) {
		try {
			BufferedReader br = new BufferedReader(new FileReader(am));
			String s = null;
			while ((s = br.readLine()) != null) {
				String[] sa = s.split(" ");
				androidMap.put(sa[0], "\"" + sa[1] + "\"");
			}
			br.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public HashSet<String> getAcquiredAppPermission() {
		if (permission_app_acquired == null) {
			readAcquiredPermission(rootPath + "/AndroidManifest.xml");
		}
		return permission_app_acquired;
	}
	
	public void checkAppPermission(String path) {
		registerMapsForClasses = new HashMap<String, HashMap<Long, List<HashMap<String, String>>>>();
		filePos2Line = new HashMap<Long, Long>();
		jumpMap = new HashMap<Long, List<Long>>();
		extendMap = new HashMap<String, String>();
		reExtendMap = new HashMap<String, List<String>>();
		
		System.out.println("decompressing apk file...");
		//String rootPath = path.split("[.]")[0];
		rootPath = decompressApp(path);
		
		if (isInvokeCompleted) {
			ic = new InvokeCompleted(rootPath + "/dedexer");
		}
		
		System.out.println("read permission from XML...");
		readAcquiredPermission(rootPath + "/AndroidManifest.xml");
		
		//deleteFolder(rootPath + "/dedexer");
		
		System.out.println("analysing classes.dex...\n");
		Dedexer dedexer = new Dedexer();
		String[] args = { "-r", "-d", rootPath + "/dedexer", rootPath + "/classes.dex" };
		long start = System.currentTimeMillis();
		dedexer.run(args);
		System.out.println("@@@ Time: " + (System.currentTimeMillis() - start) + " @@@");
		
		scanAllFile(rootPath + "/dedexer");
		
		HashSet<String> more = (HashSet<String>) permission_app_acquired.clone();
		HashSet<String> need = (HashSet<String>) permission_app_required.clone();
		comparePermission(more, need);
		showCheckResult(more, need);
		isPermissionCheck = false;
	}

	private void deleteFolder(String path) {
		File f = new File(path);
		if (f.exists()) {
			File[] fl = f.listFiles();
			for (int i = 0; i < fl.length; i++) {
				if (fl[i].isDirectory()) {
					deleteFolder(fl[i].getAbsolutePath());
				} else {
					fl[i].delete();
				}
			}
			f.delete();
		}
	}
	
	public static HashMap<String,String> getSubMap(String methodName, Long pos, int i) {
		if (registerMapsForClasses != null) {
			return registerMapsForClasses.get(methodName).get(pos).get(i);
		}
		return null;
	}
	
	private void scanAllFile(String path) {
		permission_app_required = new HashSet<String>();
		
		LinkedList list = new LinkedList();
        list.add(new File(path));
        File tmp, file[];
        while (!list.isEmpty()) {
            tmp = (File) list.removeFirst();
            if (tmp.isDirectory()) {
                file = tmp.listFiles();
                if (file == null)
                    continue;
                for (int i = 0; i < file.length; i++) {
                    if (file[i].isDirectory())
                        list.add(file[i]);
                    else {
                        analyseDDXFile(file[i].getAbsolutePath());
                    }
                }
            } else {
                analyseDDXFile(tmp.getAbsolutePath());
            }
        }
        
        apipc.checkPermission(this);
		intpc.checkPermission(this);
		cppc.checkPermission(this);
	}
	
	private String decompressApp(String path) {
		String[] ts = path.split("/");
		StringBuffer sb = new StringBuffer();
		int i;
		for (i = 0; i < ts.length - 1; i++) {
			sb.append(ts[i] + "/");
		}
		sb.append(ts[i].substring(0, ts[i].lastIndexOf(".")) + "/");
		
		File f = new File(sb.toString());
		Process process;  
        try {
			process = Runtime.getRuntime().exec(new String[]{"java", "-jar", "app/apktool.jar", "d", "-f", path, sb.toString()},null, null);
			process.waitFor();

			process = Runtime.getRuntime().exec(new String[]{"jar", "xf", path, "classes.dex"},null, new File(sb.toString()));

			process.waitFor();
			
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
        
        return sb.toString();
	}

	private void showCheckResult(HashSet<String> more, HashSet<String> need) {
		System.out.println("\nOverprivilege: ");
		if (more.size() == 0) {
			System.out.println("None");
		} else {
			for (String s : more) {
				System.out.println(s);
			}
		}
		
		System.out.println("\nRequired Permissions: ");
		if (need.size() == 0) {
			System.out.println("None");
		} else {
			for (String s : need) {
				System.out.println(s);
			}
		}
		
		System.out.println("\n\n");
	}
	
	private void comparePermission(HashSet<String> more, HashSet<String> need) {
		HashSet<String> common = new HashSet<String>();
		for (String p : more) {
			if (need.contains(p)) {
				common.add(p);
			}
		}
		for (String p : common) {
			more.remove(p);
			need.remove(p);
		}
		common.clear();
		Iterator it = need.iterator();
		while (it.hasNext()) {
			boolean flag = false;
			String[] ps = ((String) it.next()).split(" or ");
			if (ps != null) {
				for (String p : ps) {
					if (more.contains(p)) {
						common.add(p);
						flag = true;
					}
					if (p.equals("NONE")) {
						flag = true;
					}
				}
				if (flag) {
					it.remove();
				}
			}
		}
		for (String p : common) {
			more.remove(p);
		}
	}

	private void readAcquiredPermission(String path) {
		permission_app_acquired = new HashSet<String>();
		DocumentBuilderFactory domfac = DocumentBuilderFactory.newInstance();
		try {
			DocumentBuilder dombuilder = domfac.newDocumentBuilder();
			InputStream is = new FileInputStream(path);
			Document doc = dombuilder.parse(is);
			NodeList nodeList = doc.getElementsByTagName("uses-permission");
			for (int i = 0; i < nodeList.getLength(); i++) {
				Node n = nodeList.item(i);
				NamedNodeMap attributes = n.getAttributes();
				String ts = attributes.getNamedItem("android:name").toString();
				int start = ts.indexOf("\"") + 1, end = ts.indexOf("\"", start);
				PermissionChecker.setPermissionSet(permission_app_acquired, ts.substring(start, end));
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public void analyseDDXFile(String path) {
		try {
			BufferedReader br = new BufferedReader(new FileReader(path));
			String s = null;
			long lineNumber = -1, FilePos = -1;
			String methodName = null;
			String className = null;
			MethodNode cnode = null;
			while ((s = br.readLine()) != null) {
				if (s.startsWith(".class")) {
					String[] ccn = s.split(" ");
		            className = ccn[ccn.length - 1].replace('/', '.');
		            
				} else if (s.startsWith(".line")) {
					String ts = s.trim().substring(6);
					lineNumber = Integer.parseInt(ts);
					if (isCfgBuild && cnode != null) {
						cnode.addCode("\n");
						cnode.addCode(s);
					}
				} else if (s.startsWith(".FilePos")) {
					String ts = s.trim().substring(9);
					FilePos = Integer.parseInt(ts);					
					if (isCfgBuild && cnode != null) {
						cnode.addCode("\n");
						cnode.addCode(s);
					}
				} else if (s.startsWith(".method")) {
					totalMethod_++;
					String ts = s.trim().substring(8);
					String[] ta = ts.split(" ");
                    methodName = className + "." + ta[ta.length - 1].split("[(]")[0];
                    if (isInvokeCompleted) {
                    	invokeMap.put(methodName, new LinkedList<String>());
                    }
                    if (isCfgBuild) {
                    	cnode = MethodNode.createMethodNode(methodName, className, s, FilePos, lineNumber);
                    	cb.addNewMethod(cnode);
                    }
				} else if (s.startsWith(".end method")) {
					if (isInvokeCompleted && permission_func.size() != 0) {
						permission_func_required.put(methodName, (HashSet<String>) permission_func.clone());
						permission_func.clear();
					}
					if (isCfgBuild && cnode != null) {
						cnode.addCode("\n");
						cnode.addCode(s);
						cnode = null;
					}
				} else if (s.contains("invoke-")){
					if (isCfgBuild) {
						cnode.addCode("\n");
						cnode.addCode(s);
					}
					
					String[] tss = s.split(";")[0].trim().split("\t");
					String opcode = tss[0];
					tss = tss[1].split("}");
					
					String func = tss[1].substring(1).replace('/', '.');
					// check Api
					if (func.contains("java.lang.reflect.Method.invoke")) { //reflection
						String i = tss[0].substring(1).split(",")[0].substring(1);
						int reg = Integer.parseInt(i);
						apipc.setAppPermission(methodName + "@" + FilePos + "@" + reg + "@" + func + "@" + path);
						
						if (isCfgBuild) {
							String[] ret = getRegValueToCheck(methodName, FilePos, new int[]{reg}, func, path, 0, null);
							if (ret != null) {
								cb.addCfgEdge(methodName, ret[0].split(":")[1]);
							}
						}
						
					} else {						
						apipc.setAppPermission(func);
						apipc.analyse(permission_app_required, func, null);
						if (isInvokeCompleted) {
							apipc.analyse(permission_func, func, null);

							LinkedList<String> list = (LinkedList<String>) invokeMap.get(methodName);
							list.add(func);
							
							if (reInvokeMap.containsKey(func)) {
								list = (LinkedList<String>) reInvokeMap.get(func);
							} else {
								list = new LinkedList<String>();
							}
							list.add(methodName);
							reInvokeMap.put(func, list);
							
							if (func.contains("java.lang.Thread.start")) {
								getRegValueToCheck(methodName, FilePos, new int[]{0}, func, path, 0, ic);
							}
						}
						if (isCfgBuild) {
							if (func.contains("java.lang.Thread.start")) {
								String[] ret = getRegValueToCheck(methodName, FilePos, new int[]{0}, func, path, 0, null);
								String cl = ret[0].substring("Thread(".length(), ret[0].length() - 1);
								String m = cl + ".run";
								cb.addCfgEdge(methodName, m);
							} else {
								cb.addCfgEdge(methodName, func);
							}
						}
					}					

					// check Content
					// TODO: check super class
					if ((func.endsWith("ContentResolver.query") || func.endsWith("ContentResolver.insert") || func.endsWith("ContentResolver.update") || func.endsWith("ContentResolver.delete")) &&
						 !func.startsWith("java")) {
						String i;
						int reg;
						if (opcode.contains("/range")) {  // {v0..vn}
							i = tss[0].substring(1).split("[.]")[0].substring(1);
							reg = Integer.parseInt(i);
							if (!opcode.equals("invoke-static")) {
								reg++;
							}
						} else { // {va,vb,vc}
							if (!opcode.equals("invoke-static")) {
								i = tss[0].substring(1).split(",")[1].substring(1);
							} else {
								i = tss[0].substring(1).split(",")[0].substring(1);
							}
							reg = Integer.parseInt(i);
						}
						getRegValueToCheck(methodName, FilePos, new int[]{reg}, func, path, 0, cppc);
						//cppc.setAppPermission(methodName + "@" + FilePos + "@" + reg + "@" + func + "@" + path);
					}
					// check Intent
					tss = tss[0].substring(1).split("[,]");
					if (func.contains(".registerReceiver")) {
						String cl = func.substring(0, func.lastIndexOf("."));
						String fc = extendMap.get(cl);
						if ((cl.equals("android.content.Context") || cl.equals("android.app.Activity")) || 
							(fc != null && (fc.equals("android.content.Context") || fc.equals("android.app.Activity")))) {
							String i = tss[2].substring(1);
							Integer reg = Integer.parseInt(i);
							
							//intpc.setAppPermission(methodName + "@" + FilePos + "@" + reg + "@" + func + "@" + path);
							getRegValueToCheck(methodName, FilePos, new int[]{reg}, func, path, 0, intpc);
							if (isIntentCheck) {
								if (tss.length == 3) {	// registerReceiver (BroadcastReceiver receiver, IntentFilter filter)
									intc.addIntentPair(methodName, FilePos, new String[]{tss[2].substring(1)}, "receive");
								} else if (tss.length == 5) {	// registerReceiver (BroadcastReceiver receiver, IntentFilter filter, String broadcastPermission, Handler scheduler)
									intc.addIntentPair(methodName, FilePos, new String[]{tss[2].substring(1), tss[3].substring(1)}, "receive");
								}
							}
							
							if (isInvokeCompleted) { // start invoke check
								getRegValueToCheck(methodName, FilePos, new int[]{Integer.parseInt(tss[1].substring(1)), Integer.parseInt(tss[2].substring(1))}, func, path, 0, ic);
							}
						}					
					} else if (func.contains("sendBroadcast") ||
							   func.contains("sendOrderedBroadcast") ||
							   func.contains("sendStickyBroadcast") ||
							   func.contains("sendStickyOrderedBroadcast") ||
							   func.contains("startActivity") ||
							   func.contains("startActivityForResult") ||
							   func.contains("startService") ||
							   func.contains("bindService")) {
						
						String i = tss[1].substring(1);
						Integer reg = Integer.parseInt(i);
						
						intpc.setAppPermission(methodName + "@" + FilePos + "@" + reg + "@" + func + "@" + path);
						
						if (isIntentCheck) {
							if (func.contains("sendBroadcast") || func.contains("sendOrderedBroadcast")) {
								// check intent with uri & permission
								intc.addIntentPair(methodName, FilePos, new String[]{tss[1].substring(1), tss[2].substring(1)}, "send");
							}
						}
						if (isInvokeCompleted) {
							if (func.contains("startActivity") || func.contains("startActivityForResult")) {
								// TODO
							}
							if (func.contains("startService")) {
								// TODO
							}
							getRegValueToCheck(methodName, FilePos, new int[]{Integer.parseInt(tss[1].substring(1))}, func, path, 0, ic);
						}
					}
				} else {
					if (cnode != null) {
						cnode.addCode("\n");
						cnode.addCode(s);
					}
				}
				//if (invokeMap.containsKey("java.lang.reflect.Method.invoke")) {
				//	System.out.println();
				//}
			}
			br.close();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public String[] getRegValueToCheck(String methodName, long FilePos, int[] reg, String func, String path, int kind, PermissionChecker pc) {
		List<HashMap<String, String>> paramList = null;
		String[] ret = null;
		HashMap<Long, List<HashMap<String, String>>> ll = registerMapsForClasses.get(methodName);
		if (ll != null) {
			paramList = ll.get(new Long(FilePos));
			boolean flag = false;
			if (paramList != null) {
				for (HashMap<String, String> is : paramList) {
					ret = new String[reg.length];
					for (int i = 0; i < reg.length; i++) {
						String rv = is.get(String.valueOf(reg[i]));
						if (rv != null) {
							if (rv.contains("VALUE(")) { // other class's field
								int tmpi = rv.indexOf("VALUE(");
								int tmpj = rv.indexOf(")", tmpi);
								String field = rv.substring(tmpi + 6, tmpj);
								String oclassName = field.substring(0, field.lastIndexOf("."));
								HashMap<Long, List<HashMap<String, String>>> ohm = registerMapsForClasses.get(oclassName + ".<clinit>");
								if (ohm != null && !ohm.equals(ll)) { // avoid endless loop
									Iterator it = ohm.entrySet().iterator();
									while (it.hasNext()) {
										Map.Entry e1 = (Entry) it.next();
										List<HashMap<String, String>> regMapList = (List) e1.getValue();
										for (HashMap<String, String> rm : regMapList) {
											String fv = rm.get(field);
											if (fv != null && fv.contains("Uri(") && rv.contains("Uri(")) {
												fv = fv.substring(fv.indexOf("Uri(") + 5, fv.lastIndexOf(")") - 1);
											}
											ret[i] = rv.substring(0, tmpi) + fv + rv.substring(tmpj + 1);
											if (fv != null && !fv.equals("")) {
												flag = true;
												break;
											}
										}
										if (flag)	break;
									}
								}
							} else {
								ret[i] = rv;
								flag = true;
							}
						}
					}
					if (pc != null) {
						if (pc instanceof APIPermissionChecker) {
							pc.analyse(permission_app_required, methodName, ret);
						} else {
							//if (ret == null || ret[0] == null) {
							//	System.out.println();
							//}
							pc.analyse(permission_app_required, func, ret);
						}
						if (isInvokeCompleted) {
							if (func.contains("java.lang.Thread.start") && ret[0] != null && ret[0].contains("Thread(")) {
								String cl = ret[0].substring("Thread(".length(), ret[0].length() - 1).replace("/", ".");
								String m = cl + ".run";
								List list = invokeMap.get(methodName);
								list.add(m);
								
								if (reInvokeMap.containsKey(m)) {
									list = reInvokeMap.get(m);
								} else {
									list = new LinkedList<String>();
								}
								list.add(methodName);
								reInvokeMap.put(m, list);
							} else {
								if (pc instanceof APIPermissionChecker && ret != null) {
									pc.analyse(permission_func, methodName, ret);
								} else {
									pc.analyse(permission_func, func, ret);
								}
							}
						}
					}
				}
			}
			if (flag == false) {
			}
		}
		return ret;
	}
	
	public static HashSet<String> getPermissionForMethod(String m) {
		if (isInvokeCompleted && permission_func_required != null) {
			return permission_func_required.get(m);
		}
		return null;
	}
	
	public static void printRegisterMapsForClasses() {
		if (registerMapsForClasses != null && registerMapsForClasses.size() > 0) {
			Iterator i = registerMapsForClasses.entrySet().iterator();
			while (i.hasNext()) {
				Map.Entry entry = (Map.Entry) i.next();
				
				if (!entry.getKey().equals("android.exp.RegMapExp.Exp")) {
					continue;
				}
				System.out.println("@ " + entry.getKey() + " @");
				
				HashMap<Long, List<HashMap<Integer, String>>> lis = (HashMap<Long, List<HashMap<Integer, String>>>) entry
						.getValue();
				Iterator j = lis.entrySet().iterator();
				while (j.hasNext()) {
					Map.Entry e = (Map.Entry) j.next();
					long ln = filePos2Line.get(e.getKey());
					System.out.println("\t" + e.getKey());
					List<HashMap<Integer, String>> is = (List<HashMap<Integer, String>>) e
							.getValue();
					for (HashMap<Integer, String> it : is) {
						Iterator k = it.entrySet().iterator();
						while (k.hasNext()) {
							Map.Entry ee = (Map.Entry) k.next();
							System.out.println("\t\t" + ee.getKey() + ": "
									+ ee.getValue());
						}
						System.out.println("------------------------------");
					}
				}
			}
		} else {
			System.out.println("registerMapsForClasses is null or empty");
		}
	}

	public void printPermissionForMethod() {
		if (permission_func_required != null) {
			Iterator i = permission_func_required.entrySet().iterator();
			while (i.hasNext()) {
				Map.Entry e = (Entry) i.next();
				System.out.println("Method: " + e.getKey().toString());
				
				for (String s : (HashSet<String>)e.getValue()) {
					System.out.println("\t" + s);
				}
			}
		}
	}
	
	private void printInvokeMap(HashMap<String, List<String>> map) {
		if (map != null) {
			Iterator i = map.entrySet().iterator();
			while (i.hasNext()) {
				Map.Entry e = (Entry)i.next();
				System.out.println("Method: " + e.getKey().toString());
				
				for (String s : (LinkedList<String>)e.getValue()) {
					System.out.println("\t\t" + s);
				}
			}
		}
	}
	
	public static void main(String[] args) {
		long startt = System.currentTimeMillis();
		if (args.length < 5) {
			System.out.println("Missing arg...");
		} else {
			DedexerPlus dp = new DedexerPlus(args[0], args[1], args[2], args[3]);
			try {
				JasminStyleCodeGenerator.setTimeLimit(Long.parseLong(args[5]));
				dp.checkAppPermission(args[4]);
			} catch (Exception e) {
				e.printStackTrace();
			}
			System.out.println("Total: " + totalMethod);
			//dp.printPermissionForMethod();
			//dp.printInvokeMap(invokeMap);
			//dp.printInvokeMap(reInvokeMap);
			//dp.ic.checkPermission(null);
			//((InvokeCompleted)dp.ic).initEndsFromXML("D:\\Temp\\Mms\\AndroidManifest.xml");
			//((InvokeCompleted)dp.ic).findPermissionLeakAction();
			//dp.cb.printMethodMap();
			//((InvokeCompleted)dp.ic).findDeadMethod(invokeMap, reInvokeMap);
			System.out.println("Finish: " + totalMethod_);
		}
		//printRegisterMapsForClasses();
		//System.out.println("### " + (System.currentTimeMillis() - startt) + " ###");
	}
}
